import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import cv2
import pytesseract
from pytesseract import Output 
import numpy as np
from openpyxl import Workbook

class Application(tk.Frame):
    def __init__(self, master=None):
        super().__init__(master)
        self.master = master
        self.master.title("图片表格转Excel工具")
        self.pack()
        self.create_widgets()
        
    def create_widgets(self):
        # 输入文件选择
        self.input_label = tk.Label(self, text="输入图片:")
        self.input_label.grid(row=0, column=0, sticky="w")
        
        self.input_entry = tk.Entry(self, width=40)
        self.input_entry.grid(row=0, column=1)
        
        self.input_button = tk.Button(self, text="选择...", command=self.select_input)
        self.input_button.grid(row=0, column=2)
        
        # 输出文件选择
        self.output_label = tk.Label(self, text="输出Excel:")
        self.output_label.grid(row=1, column=0, sticky="w")
        
        self.output_entry = tk.Entry(self, width=40)
        self.output_entry.grid(row=1, column=1)
        
        self.output_button = tk.Button(self, text="选择...", command=self.select_output)
        self.output_button.grid(row=1, column=2)
        
        # 执行按钮
        self.run_button = tk.Button(self, text="开始转换", command=self.run_conversion)
        self.run_button.grid(row=2, column=1, pady=10)
        
        # 状态显示
        self.status = tk.StringVar()
        self.status.set("准备就绪")
        self.status_label = tk.Label(self, textvariable=self.status)
        self.status_label.grid(row=3, column=0, columnspan=3)
        
        # 进度条
        self.progress = ttk.Progressbar(self, orient="horizontal", length=200, mode="determinate")
        self.progress.grid(row=4, column=0, columnspan=3, pady=5)
    
    def select_input(self):
        filename = filedialog.askopenfilename(
            title="选择输入图片",
            filetypes=[("图片文件", "*.png;*.jpg;*.jpeg;*.bmp"), ("所有文件", "*.*")]
        )
        if filename:
            self.input_entry.delete(0, tk.END)
            self.input_entry.insert(0, filename)
    
    def select_output(self):
        filename = filedialog.asksaveasfilename(
            title="保存Excel文件",
            defaultextension=".xlsx",
            filetypes=[("Excel文件", "*.xlsx"), ("所有文件", "*.*")]
        )
        if filename:
            self.output_entry.delete(0, tk.END)
            self.output_entry.insert(0, filename)
    
    def run_conversion(self):
        input_file = self.input_entry.get()
        output_file = self.output_entry.get()
        
        if not input_file or not output_file:
            messagebox.showerror("错误", "请选择输入和输出文件")
            return
            
        self.status.set("正在处理...")
        self.progress["value"] = 30
        self.update()
        
        try:
            enhance_table_recognition(input_file, output_file)
            self.status.set(f"转换完成: {output_file}")
            self.progress["value"] = 100
            messagebox.showinfo("成功", "表格转换完成!")
        except Exception as e:
            self.status.set(f"错误: {str(e)}")
            self.progress["value"] = 0
            messagebox.showerror("错误", f"处理失败: {str(e)}")

def enhance_table_recognition(image_path, output_excel):
    try:
        # 读取并预处理图像
        img = cv2.imread(image_path)
        if img is None:
            raise FileNotFoundError(f"无法读取图像: {image_path}")
        
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # 增强对比度和降噪
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        enhanced = clahe.apply(gray)
        
        # 检测表格线
        thresh = cv2.threshold(enhanced, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
        detected_lines = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)
        
        # 使用表格专用OCR配置
        custom_config = r'--oem 3 --psm 6 -l chi_sim+eng'
        data = pytesseract.image_to_data(enhanced, config=custom_config, output_type=Output.DICT)
        
        # 创建Excel并填充数据
        wb = Workbook()
        ws = wb.active
        
        # 按单元格位置组织数据
        cell_data = {}
        for i in range(len(data['text'])):
            if int(data['conf'][i]) > 60:  # 只保留置信度高的结果
                x, y = data['left'][i], data['top'][i]
                cell_data[(x,y)] = data['text'][i]
        
        # 按坐标排序后写入Excel
        for idx, (pos, text) in enumerate(sorted(cell_data.items())):
            row = pos[1] // 30  # 估算行号(根据实际调整)
            col = pos[0] // 100  # 估算列号
            ws.cell(row=row+1, column=col+1, value=text)
        
        wb.save(output_excel)
        print(f"表格已保存至: {output_excel}")

    except FileNotFoundError as e:
        print(f"文件错误: {str(e)}")
    except ValueError as e:
        print(f"值错误: {str(e)}")
    except Exception as e:
        print(f"处理失败: {str(e)}")
        # 保存中间处理结果用于调试
        cv2.imwrite("debug_processed.png", enhanced)
        cv2.imwrite("debug_detected_lines.png", detected_lines*255)


def preprocess_cell(cell_img):
    """预处理单元格图像"""
    if len(cell_img) == 0 or len(cell_img[0]) == 0:
        return np.zeros((10, 10), dtype=np.uint8) + 255
    
    gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    if np.mean(thresh) < 127:
        thresh = cv2.bitwise_not(thresh)
    
    kernel = np.ones((1, 1), np.uint8)
    processed = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    processed = cv2.copyMakeBorder(processed, 15, 15, 15, 15, cv2.BORDER_CONSTANT, value=255)
    return processed


if __name__ == "__main__":
    root = tk.Tk()
    app = Application(master=root)
    app.mainloop()
